import copy
import inspect
import warnings
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union

import torch
import torch.distributed as dist
from torch import nn
import copy

from transformers.generation.logits_process import (
    LogitsProcessorList,
)
from transformers.generation.stopping_criteria import (
    StoppingCriteria,
    StoppingCriteriaList,
    validate_stopping_criteria,
)
import transformers
from transformers.generation.utils import SampleOutput, GenerateEncoderDecoderOutput, GenerateDecoderOnlyOutput, GenerateNonBeamOutput 
from transformers.generation.utils import GenerationMixin

def greedy_search(
        self,
        input_ids: torch.LongTensor,
        logits_processor: Optional[LogitsProcessorList] = None,
        stopping_criteria: Optional[StoppingCriteriaList] = None,
        max_length: Optional[int] = None,
        pad_token_id: Optional[int] = None,
        eos_token_id: Optional[Union[int, List[int]]] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_scores: Optional[bool] = None,
        return_dict_in_generate: Optional[bool] = None,
        synced_gpus: bool = False,
        streamer: Optional["BaseStreamer"] = None,
        **model_kwargs,
    ) -> Union[GenerateNonBeamOutput, torch.LongTensor]:
        r"""
        Generates sequences of token ids for models with a language modeling head using **greedy decoding** and can be
        used for text-decoder, text-to-text, speech-to-text, and vision-to-text models.

        <Tip warning={true}>

        In most cases, you do not need to call [`~generation.GenerationMixin.greedy_search`] directly. Use generate()
        instead. For an overview of generation strategies and code examples, check the [following
        guide](../generation_strategies).

        </Tip>


        Parameters:
            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
                The sequence used as a prompt for the generation.
            logits_processor (`LogitsProcessorList`, *optional*):
                An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`]
                used to modify the prediction scores of the language modeling head applied at each generation step.
            stopping_criteria (`StoppingCriteriaList`, *optional*):
                An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`]
                used to tell if the generation loop should stop.

            max_length (`int`, *optional*, defaults to 20):
                **DEPRECATED**. Use `logits_processor` or `stopping_criteria` directly to cap the number of generated
                tokens. The maximum length of the sequence to be generated.
            pad_token_id (`int`, *optional*):
                The id of the *padding* token.
            eos_token_id (`Union[int, List[int]]`, *optional*):
                The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens.
            output_attentions (`bool`, *optional*, defaults to `False`):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more details.
            output_hidden_states (`bool`, *optional*, defaults to `False`):
                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
                for more details.
            output_scores (`bool`, *optional*, defaults to `False`):
                Whether or not to return the prediction scores. See `scores` under returned tensors for more details.
            return_dict_in_generate (`bool`, *optional*, defaults to `False`):
                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
            synced_gpus (`bool`, *optional*, defaults to `False`):
                Whether to continue running the while loop until max_length (needed for ZeRO stage 3)
            streamer (`BaseStreamer`, *optional*):
                Streamer object that will be used to stream the generated sequences. Generated tokens are passed
                through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
            model_kwargs:
                Additional model specific keyword arguments will be forwarded to the `forward` function of the model.
                If model is an encoder-decoder model the kwargs should include `encoder_outputs`.

        Return:
            [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`] or
            `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a
            [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and
            `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if
            `model.config.is_encoder_decoder=True`.

        Examples:

        ```python
        >>> from transformers import (
        ...     AutoTokenizer,
        ...     AutoModelForCausalLM,
        ...     LogitsProcessorList,
        ...     MinLengthLogitsProcessor,
        ...     StoppingCriteriaList,
        ...     MaxLengthCriteria,
        ... )

        >>> tokenizer = AutoTokenizer.from_pretrained("gpt2")
        >>> model = AutoModelForCausalLM.from_pretrained("gpt2")

        >>> # set pad_token_id to eos_token_id because GPT2 does not have a PAD token
        >>> model.generation_config.pad_token_id = model.generation_config.eos_token_id

        >>> input_prompt = "It might be possible to"
        >>> input_ids = tokenizer(input_prompt, return_tensors="pt").input_ids

        >>> # instantiate logits processors
        >>> logits_processor = LogitsProcessorList(
        ...     [
        ...         MinLengthLogitsProcessor(10, eos_token_id=model.generation_config.eos_token_id),
        ...     ]
        ... )
        >>> stopping_criteria = StoppingCriteriaList([MaxLengthCriteria(max_length=20)])

        >>> outputs = model.greedy_search(
        ...     input_ids, logits_processor=logits_processor, stopping_criteria=stopping_criteria
        ... )

        >>> tokenizer.batch_decode(outputs, skip_special_tokens=True)
        ["It might be possible to get a better understanding of the nature of the problem, but it's not"]
        ```"""
        # init values
        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
        if max_length is not None:
            warnings.warn(
                "`max_length` is deprecated in this function, use"
                " `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_length=max_length)])` instead.",
                UserWarning,
            )
            stopping_criteria = validate_stopping_criteria(stopping_criteria, max_length)
        pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id
        eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id
        if isinstance(eos_token_id, int):
            eos_token_id = [eos_token_id]
        eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None
        output_scores = output_scores if output_scores is not None else self.generation_config.output_scores
        output_attentions = (
            output_attentions if output_attentions is not None else self.generation_config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states
        )

        return_dict_in_generate = (
            return_dict_in_generate
            if return_dict_in_generate is not None
            else self.generation_config.return_dict_in_generate
        )

        # init attention / hidden states / scores tuples
        scores = () if (return_dict_in_generate and output_scores) else None
        decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
        cross_attentions = () if (return_dict_in_generate and output_attentions) else None
        decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None

        # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
        if return_dict_in_generate and self.config.is_encoder_decoder:
            encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
            encoder_hidden_states = (
                model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
            )

        # keep track of which sequences are already finished
        unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)

        this_peer_finished = False  # used by synced_gpus only
            

        use_cf = model_kwargs.get("use_cf")
        if use_cf:
            output_attentions_cf = (
            output_attentions if output_attentions is not None else self.generation_config.output_attentions
            )
            output_hidden_states_cf = (
                output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states
            )
            model_kwargs_cf = copy.deepcopy(model_kwargs)
            context_free_indices = model_kwargs_cf["context_free_indices"]
            if len(context_free_indices) > 0:
                for i, indices in enumerate(context_free_indices):
                    # exclude_indices = torch.arange(indices[0], indices[1], device=input_ids.device)
                    # reserved_indices = torch.arange(0, model_kwargs_cf["inputs_embeds"][i].shape[0], device=input_ids.device)
                    # mask = torch.ones_like(reserved_indices)
                    # mask[exclude_indices] = 0
                    # reserved_indices = reserved_indices[mask == 1]
                    # if model_kwargs_cf.get("inputs_embeds") is not None:
                    #     model_kwargs_cf["inputs_embeds"] = model_kwargs_cf["inputs_embeds"][i,reserved_indices].unsqueeze(0)
                    # if model_kwargs_cf.get("position_ids") is not None:
                    #     model_kwargs_cf["position_ids"] = model_kwargs_cf["position_ids"][i,reserved_indices].unsqueeze(0)
                    # if model_kwargs_cf.get("attention_mask") is not None:
                    #     model_kwargs_cf["attention_mask"] = model_kwargs_cf["attention_mask"][i,reserved_indices].unsqueeze(0)
                    # if model_kwargs_cf.get("past_key_values") is not None:
                    
                    indices_range = torch.arange(indices[0], indices[1], device=input_ids.device)
                    # Randomly select 50% for masking
                    mask = torch.rand(indices_range.shape[0], device=input_ids.device) < 0.5
                    reserved_indices = indices_range[mask]
                    model_kwargs_cf["attention_mask"][i, reserved_indices] = False
        while True:
            if synced_gpus:
                # Under synced_gpus the `forward` call must continue until all gpus complete their sequence.
                # The following logic allows an early break if all peers finished generating their sequence
                this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0).to(input_ids.device)
                # send 0.0 if we finished, 1.0 otherwise
                dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM)
                # did all peers finish? the reduced sum will be 0.0 then
                if this_peer_finished_flag.item() == 0.0:
                    break

            # prepare model inputs
            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
    
            # forward pass to get next token
            outputs = self(
                **model_inputs,
                return_dict=True,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
            )


            if synced_gpus and this_peer_finished:
                continue  # don't waste resources running the code we don't need

            next_token_logits = outputs.logits[:, -1, :]

            ## For context free decoding initial
            # use_cf = model_kwargs.get("use_cf")
            if use_cf:
                # output_attentions_cf = (
                # output_attentions if output_attentions is not None else self.generation_config.output_attentions
                # )
                # output_hidden_states_cf = (
                #     output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states
                # )
                # model_kwargs_cf = copy.deepcopy(model_kwargs)
                # context_free_indices = model_kwargs_cf["context_free_indices"]
                # if len(context_free_indices) > 0:
                #     for i, indices in enumerate(context_free_indices):
                #         # exclude_indices = torch.arange(indices[0], indices[1], device=input_ids.device)
                #         # reserved_indices = torch.arange(0, model_kwargs_cf["inputs_embeds"][i].shape[0], device=input_ids.device)
                #         # mask = torch.ones_like(reserved_indices)
                #         # mask[exclude_indices] = 0
                #         # reserved_indices = reserved_indices[mask == 1]
                #         # if model_kwargs_cf.get("inputs_embeds") is not None:
                #         #     model_kwargs_cf["inputs_embeds"] = model_kwargs_cf["inputs_embeds"][i,reserved_indices].unsqueeze(0)
                #         # if model_kwargs_cf.get("position_ids") is not None:
                #         #     model_kwargs_cf["position_ids"] = model_kwargs_cf["position_ids"][i,reserved_indices].unsqueeze(0)
                #         # if model_kwargs_cf.get("attention_mask") is not None:
                #         #     model_kwargs_cf["attention_mask"] = model_kwargs_cf["attention_mask"][i,reserved_indices].unsqueeze(0)
                #         # if model_kwargs_cf.get("past_key_values") is not None:

                #         model_kwargs_cf["attention_mask"][i,indices[0]:indices[1]] = False
                # model_kwargs.pop("position_ids_cf", None)
                # model_kwargs.pop("attention_mask_cf", None)
                # model_kwargs_cf["inputs_embeds"] = model_kwargs_cf.get("inputs_embeds_cf")
                # model_kwargs_cf["position_ids"] = model_kwargs_cf.get("position_ids_cf")
                # model_kwargs_cf["attention_mask"] = model_kwargs_cf.get("attention_mask_cf")
                # if model_kwargs_cf["attention_mask"] is None:
                #     model_kwargs_cf["attention_mask"] = self._prepare_attention_mask_for_generation(inputs=model_kwargs_cf["inputs_embeds"], pad_token_id=pad_token_id, eos_token_id=eos_token_id)
                # model_kwargs_cf.pop("inputs_embeds_cf", None)
                # model_kwargs_cf.pop("position_ids_cf", None)
                # model_kwargs_cf.pop("attention_mask_cf", None)
            
                model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs_cf)
                outputs_cf = self(
                    **model_inputs,
                    return_dict=True,
                    output_attentions=output_attentions_cf,
                    output_hidden_states=output_hidden_states_cf,
                    use_cf=use_cf,
                )
                next_token_logits_cf = outputs_cf.logits[:, -1, :]
                ## pre-process logits from context free inputs
                cf_alpha = model_kwargs_cf.get("cf_alpha") if model_kwargs_cf.get("cf_alpha") is not None else 1.0
                cf_beta = model_kwargs_cf.get("cf_beta") if model_kwargs_cf.get("cf_beta") is not None else 0.1

                cutoff = torch.log(torch.tensor(cf_beta)) + next_token_logits.max(dim=-1, keepdim=True).values
            
                diffs = (1+cf_alpha)*next_token_logits - cf_alpha*next_token_logits_cf
                cf_logits = diffs.masked_fill(next_token_logits < cutoff, -float("inf"))

                cf_logits = logits_processor(input_ids, cf_logits)
                next_tokens = torch.argmax(cf_logits, dim=-1)
            else:
                # pre-process distribution
                next_tokens_scores = logits_processor(input_ids, next_token_logits)
                # argmax
                next_tokens = torch.argmax(next_tokens_scores, dim=-1)

            # Store scores, attentions and hidden_states when required
            if return_dict_in_generate:
                if output_scores:
                    scores += (next_tokens_scores,)
                if output_attentions:
                    decoder_attentions += (
                        (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)
                    )
                    if self.config.is_encoder_decoder:
                        cross_attentions += (outputs.cross_attentions,)

                if output_hidden_states:
                    decoder_hidden_states += (
                        (outputs.decoder_hidden_states,)
                        if self.config.is_encoder_decoder
                        else (outputs.hidden_states,)
                    )

            # finished sentences should have their next token be a padding token
            if eos_token_id is not None:
                if pad_token_id is None:
                    raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.")
                next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)

            # update generated ids, model inputs, and length for next step
            input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
            if streamer is not None:
                streamer.put(next_tokens.cpu())
            model_kwargs = self._update_model_kwargs_for_generation(
                outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
            )
            ## update model_kwargs_cf for context free decoding
            if use_cf:
                model_kwargs_cf = self._update_model_kwargs_for_generation(
                    outputs_cf, model_kwargs_cf, is_encoder_decoder=self.config.is_encoder_decoder
                )
            # if eos_token was found in one sentence, set sentence to finished
            if eos_token_id_tensor is not None:
                unfinished_sequences = unfinished_sequences.mul(
                    next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
                )

                # stop when each sentence is finished
                if unfinished_sequences.max() == 0:
                    this_peer_finished = True

            # stop if we exceed the maximum length
            if stopping_criteria(input_ids, scores):
                this_peer_finished = True

            if this_peer_finished and not synced_gpus:
                break

        if streamer is not None:
            streamer.end()

        if return_dict_in_generate:
            if self.config.is_encoder_decoder:
                return GenerateEncoderDecoderOutput(
                    sequences=input_ids,
                    scores=scores,
                    encoder_attentions=encoder_attentions,
                    encoder_hidden_states=encoder_hidden_states,
                    decoder_attentions=decoder_attentions,
                    cross_attentions=cross_attentions,
                    decoder_hidden_states=decoder_hidden_states,
                    past_key_values=model_kwargs.get("past_key_values"),
                )
            else:
                return GenerateDecoderOnlyOutput(
                    sequences=input_ids,
                    scores=scores,
                    attentions=decoder_attentions,
                    hidden_states=decoder_hidden_states,
                    past_key_values=model_kwargs.get("past_key_values"),
                )
        else:
            return input_ids

def evolve_cf_sampling():
    transformers.generation.utils.GenerationMixin.greedy_search = greedy_search